from django.db import models
from django.contrib.auth import get_user_model
from django.utils.translation import gettext_lazy as _
from data_analysis.models import DataSource

User = get_user_model()

class CrawlerTask(models.Model):
    """爬虫任务模型"""
    class TaskStatus(models.TextChoices):
        PENDING = 'pending', _('等待中')
        RUNNING = 'running', _('运行中')
        COMPLETED = 'completed', _('已完成')
        FAILED = 'failed', _('失败')
        CANCELLED = 'cancelled', _('已取消')

    name = models.CharField(_('任务名称'), max_length=100)
    data_source = models.ForeignKey(DataSource, on_delete=models.CASCADE, related_name='crawler_tasks')
    status = models.CharField(
        _('状态'),
        max_length=20,
        choices=TaskStatus.choices,
        default=TaskStatus.PENDING
    )
    parameters = models.JSONField(_('爬取参数'), default=dict)
    progress = models.IntegerField(_('进度'), default=0)
    error_message = models.TextField(_('错误信息'), blank=True)
    total_items = models.IntegerField(_('总条目数'), default=0)
    crawled_items = models.IntegerField(_('已爬取条目数'), default=0)
    
    user = models.ForeignKey(User, on_delete=models.CASCADE, related_name='crawler_tasks')
    created_at = models.DateTimeField(_('创建时间'), auto_now_add=True)
    updated_at = models.DateTimeField(_('更新时间'), auto_now=True)
    completed_at = models.DateTimeField(_('完成时间'), null=True, blank=True)

    class Meta:
        verbose_name = _('爬虫任务')
        verbose_name_plural = _('爬虫任务')
        ordering = ['-created_at']

    def __str__(self):
        return self.name

class CrawledData(models.Model):
    """爬取的数据模型"""
    task = models.ForeignKey(CrawlerTask, on_delete=models.CASCADE, related_name='crawled_data')
    data = models.JSONField(_('数据内容'))
    metadata = models.JSONField(_('元数据'), default=dict)
    created_at = models.DateTimeField(_('创建时间'), auto_now_add=True)

    class Meta:
        verbose_name = _('爬取数据')
        verbose_name_plural = _('爬取数据')
        ordering = ['-created_at']

    def __str__(self):
        return f"Data for task {self.task.name}"
